package cn.sciento.boot.api.customize.agent;

import com.sun.tools.attach.VirtualMachine;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.lang.management.ManagementFactory;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import java.util.jar.Manifest;
import javassist.CannotCompileException;
import javassist.ClassClassPath;
import javassist.ClassPath;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.NotFoundException;
import cn.sciento.boot.api.customize.exception.VMNotFoundException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;

public class HotSwapClient {
  private static final String CUSTOMIZE_AGENT_DIR = "customize" + File.separator + "agent" + File.separator;
  
  private static final Logger LOGGER = LoggerFactory.getLogger(HotSwapClient.class);
  
  private static Instrumentation instrumentation = null;
  
  public static void setInstrumentation(Instrumentation instrumentation) {
    LOGGER.info("hot swap get Instrumentation is {}", instrumentation);
    HotSwapClient.instrumentation = instrumentation;
  }
  
  public static Instrumentation getInstrumentation() {
    return instrumentation;
  }
  
  public static void redefine(Class<?> oldClass, CtClass newClass) throws NotFoundException, IOException, CannotCompileException, VMNotFoundException {
    Class<?>[] old = new Class[] { oldClass };
    CtClass[] newClasses = { newClass };
    redefine(old, newClasses);
  }
  
  public static void redefine(Class<?>[] oldClasses, CtClass[] newClasses) throws NotFoundException, IOException, CannotCompileException, VMNotFoundException {
    startAgent();
    try {
      for (int i = 0; i < oldClasses.length; i++) {
        ClassDefinition classDefinition = new ClassDefinition(oldClasses[i], newClasses[i].toBytecode());
        instrumentation.redefineClasses(new ClassDefinition[] { classDefinition });
        LOGGER.info("redefine [{}] success ...", oldClasses[i].getName());
      } 
    } catch (ClassNotFoundException e) {
      throw new NotFoundException(e.getMessage(), e);
    } catch (UnmodifiableClassException e) {
      throw new CannotCompileException(e.getMessage(), e);
    } 
  }
  
  public static void startAgent() throws NotFoundException, VMNotFoundException {
    if (instrumentation != null)
      return; 
    try {
      Thread.currentThread().getContextClassLoader().loadClass("com.sun.tools.attach.VirtualMachine");
    } catch (ClassNotFoundException e) {
      LOGGER.warn("did not start hot swap agent, can't dynamic redefine class, please check if include tools.jar in the package.");
      throw new VMNotFoundException("VirtualMachine not found");
    } 
    try {
      File agentJar = createJarFile();
      String nameOfRunningVm = ManagementFactory.getRuntimeMXBean().getName();
      String pid = nameOfRunningVm.substring(0, nameOfRunningVm.indexOf('@'));
      VirtualMachine vm = VirtualMachine.attach(pid);
      vm.loadAgent(agentJar.getPath());
      vm.detach();
      LOGGER.info("attach to vm [{}] success. instrumentation is {}", pid, instrumentation);
    } catch (Exception e) {
      throw new NotFoundException("hotswap agent", e);
    } 
    for (int sec = 0; sec < 20; sec++) {
      if (instrumentation != null)
        return; 
      try {
        LOGGER.info("waiting for instrumentation {}s ...", Integer.valueOf(sec + 1));
        Thread.sleep(1000L);
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        break;
      } 
    } 
    throw new NotFoundException("hotswap agent (timeout)");
  }
  
  private static File createJarFile() throws IOException, CannotCompileException, NotFoundException {
    LOGGER.info("instrumentation is null, begin create agent jar ...");
    File jar = new File(CUSTOMIZE_AGENT_DIR + "customize-agent.jar");
    jar.deleteOnExit();
    boolean make = jar.getParentFile().mkdirs();
    Assert.isTrue(make, "create dir error!");
    boolean create = jar.createNewFile();
    Assert.isTrue(create, "create new file error!");
    File agentJar = createJarFile(jar);
    LOGGER.info("create agent jar success, path is {}", agentJar.getPath());
    return agentJar;
  }
  
  private static File createJarFile(File jar) throws IOException, CannotCompileException, NotFoundException {
    String agentName = HotSwapAgent.class.getName();
    Manifest manifest = new Manifest();
    Attributes attrs = manifest.getMainAttributes();
    attrs.put(Attributes.Name.MANIFEST_VERSION, "1.0");
    attrs.put(new Attributes.Name("Premain-Class"), agentName);
    attrs.put(new Attributes.Name("Agent-Class"), agentName);
    attrs.put(new Attributes.Name("Can-Retransform-Classes"), "true");
    attrs.put(new Attributes.Name("Can-Redefine-Classes"), "true");
    attrs.put(new Attributes.Name("Specification-Title"), "customize-agent");
    try (JarOutputStream jos = new JarOutputStream(new FileOutputStream(jar), manifest)) {
      ClassPool pool = ClassPool.getDefault();
      pool.appendClassPath((ClassPath)new ClassClassPath(HotSwapAgent.class));
      JarEntry e = new JarEntry(agentName.replace('.', '/') + ".class");
      jos.putNextEntry(e);
      jos.write(pool.get(agentName).toBytecode());
      jos.closeEntry();
    } 
    return jar;
  }
}
